#!/bin/bash

set -e

# 配置变量
PROJECT_DIR="/root/FaceFusion"
VENV_PATH="$PROJECT_DIR/facefusion_env/bin/activate"
API_DIR="$PROJECT_DIR/api"
HOST="0.0.0.0"
PORT=8000
WORKERS=1

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 输出函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 检查虚拟环境是否存在
check_venv() {
    if [ ! -f "$VENV_PATH" ]; then
        log_error "虚拟环境不存在: $VENV_PATH"
        exit 1
    fi
}

# 检查API目录是否存在
check_api_dir() {
    if [ ! -d "$API_DIR" ]; then
        log_error "API目录不存在: $API_DIR"
        exit 1
    fi
}

# 检查端口是否被占用
check_port() {
    if lsof -Pi :$PORT -sTCP:LISTEN -t >/dev/null ; then
        log_warning "端口 $PORT 已被占用"
        return 1
    fi
    return 0
}

# 优雅关闭处理
cleanup() {
    log_info "正在关闭 FaceFusion API 服务..."
    deactivate 2>/dev/null || true
    log_success "服务已关闭"
    exit 0
}

# 注册信号处理
trap cleanup SIGINT SIGTERM

# 主函数
main() {
    echo -e "${BLUE}"
    echo "==========================================="
    echo "  FaceFusion API 服务启动"
    echo "==========================================="
    echo -e "${NC}"
    echo "项目路径: $PROJECT_DIR"
    echo "虚拟环境: $VENV_PATH"
    echo "API端口: $PORT"
    echo "Workers: $WORKERS"
    echo -e "${BLUE}===========================================${NC}"
    echo ""

    # 前置检查
    log_info "执行前置检查..."
    check_venv
    check_api_dir
    check_port
    log_success "前置检查通过"

    # 切换到项目目录
    cd "$PROJECT_DIR"

    # 激活虚拟环境
    log_info "激活 Python 虚拟环境..."
    source "$VENV_PATH"
    log_success "虚拟环境激活成功"

    # 设置环境变量
    export PYTHONPATH="$PROJECT_DIR:$PYTHONPATH"
    export CUDA_VISIBLE_DEVICES=0

    log_success "🚀 Starting FaceFusion API..."
    log_info "📖 API文档: http://$HOST:$PORT/docs"
    log_info "📍 健康检查: http://$HOST:$PORT/health"
    echo ""

    # 启动 uvicorn 服务器
    uvicorn api.main:app \
        --host "$HOST" \
        --port "$PORT" \
        --workers "$WORKERS" \
        --log-level info
}

# 执行主函数
main
